/*
 * Copyright 2010 Bo Zhou<Bo.Schwarzstein@gmail.com>
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at

 * http://www.apache.org/licenses/LICENSE-2.0

 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "PRTTranslator.h"

#include <maya/MDagPath.h>
#include <maya/MDagPathArray.h>
#include <maya/MDoubleArray.h>
#include <maya/MFnNObjectData.h>
#include <maya/MFnPlugin.h>
#include <maya/MGlobal.h>
#include <maya/MItDag.h>
#include <maya/MObjectArray.h>
#include <maya/MFnParticleSystem.h>
#include <maya/MSelectionList.h>
#include <maya/MFnDoubleArrayData.h>
#include <maya/MFnVectorArrayData.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MFnTypedAttribute.h>
#include <maya/MPlug.h>
#include <maya/MVector.h>
#include <maya/MVectorArray.h>

#include <cassert>
#include <memory>

#include <zlib.h>

// position, velocity, radius(PP), rgbPP, opacityPP, id
const int PRTTranslator::DWORDCountOfSingleParticle = 3 + 3 + 1 + 3 + 1 + 1;

void* PRTTranslator::Creator()
{
	return new PRTTranslator;
}

PRTTranslator::PRTTranslator()
{
}

PRTTranslator::~PRTTranslator()
{
}

MStatus PRTTranslator::writer (const MFileObject& file, const MString& optionsString, FileAccessMode mode)
{
	if (mode != kExportActiveAccessMode)
	{
		return MStatus::kInvalidParameter;
	}

	return (Save(file.expandedFullName()) == true) ? MStatus::kSuccess : MStatus::kFailure;
}

bool PRTTranslator::haveWriteMethod () const
{
	return true;
}

bool PRTTranslator::Extract(boost::shared_array<unsigned char>& Data, long long& ParticleCount)
{
	try
	{
		// Get out selected node.
		MDagPathArray ParticleShapeNodePaths, NParticleShapeNodePaths;

		MSelectionList Selected;
		MGlobal::getActiveSelectionList(Selected);
		if (Selected.length() != 1)
		{
			MGlobal::displayError("You should select only one particle shape node.");
			return false;
		}

		// Extract Node's data.
		MVectorArray Position;
		MVectorArray Velocity;
		MDoubleArray Radius;
		MVectorArray Color;
		MDoubleArray Opacity;
		MIntArray Id;

		// Get DAG path from selected node, NParticle node or Particle node.
		MDagPath NodePath;
		Selected.getDagPath(0, NodePath);

		if (NodePath.hasFn(MFn::kNParticle))
		{
			NodePath.extendToShape();
			MFnDependencyNode Node(NodePath.node());

			// Position, 3 x float
			MPlug PositionPlug = Node.findPlug("position");
			MObject PositionDataObject;
			PositionPlug.getValue(PositionDataObject);
			assert(PositionDataObject.apiType() == MFn::kVectorArrayData);
			MFnVectorArrayData PositionData(PositionDataObject);
			PositionData.copyTo(Position);

			// Velocity, 3 x float
			MPlug VelocityPlug = Node.findPlug("velocity");
			MObject VelocityDataObject;
			VelocityPlug.getValue(VelocityDataObject);
			assert(VelocityDataObject.apiType() == MFn::kVectorArrayData);
			MFnVectorArrayData VelocityData(VelocityDataObject);
			VelocityData.copyTo(Velocity);

			// Radius, 1 x float
			MPlug RadiusPlug = Node.findPlug("radiusPP");
			MObject RadiusDataObject;
			RadiusPlug.getValue(RadiusDataObject);

			if (RadiusDataObject.apiType() == MFn::kInvalid)
			{
				RadiusPlug = Node.findPlug("radius");
				Radius.setLength(Position.length());
				for (unsigned int i = 0; i < Radius.length(); ++ i)
				{
					Radius[i] = RadiusPlug.asDouble();
				}
			}else
			{
				assert(RadiusDataObject.apiType() == MFn::kDoubleArrayData);
				MFnDoubleArrayData RadiusData(RadiusDataObject);
				RadiusData.copyTo(Radius);
			}

			// Color, 3 x float
			MPlug ColorPlug = Node.findPlug("rgbPP");
			MObject ColorDataObject;
			ColorPlug.getValue(ColorDataObject);
			if (ColorDataObject.apiType() == MFn::kInvalid)
			{
				MPlug ConstantColorRedPlug = Node.findPlug("colorRed");
				MPlug ConstantColorBluePlug = Node.findPlug("colorBlue");
				MPlug ConstantColorGreenPlug = Node.findPlug("colorGreen");

				Color.setLength(Position.length());
				for (unsigned int i = 0; i < Color.length(); ++ i)
				{
					Color[i].x = ConstantColorRedPlug.asDouble();
					Color[i].y = ConstantColorBluePlug.asDouble();
					Color[i].z = ConstantColorGreenPlug.asDouble();
				}
			}
			else
			{
				MFnVectorArrayData ColorData(ColorDataObject);
				ColorData.copyTo(Color);
			}

			// Opacity, 1 x float
			MPlug OpacityPlug = Node.findPlug("opacityPP");
			MObject OpacityDataObject;
			OpacityPlug.getValue(OpacityDataObject);
			if (OpacityDataObject.apiType() == MFn::kInvalid)
			{
				MPlug ConstantOpacityPlug = Node.findPlug("opacity");
				Opacity.setLength(Position.length());
				for (unsigned int i = 0; i < Opacity.length(); ++ i)
				{
					Opacity[i] = ConstantOpacityPlug.asDouble();
				}
			}
			else
			{
				MFnDoubleArrayData OpacityData(OpacityDataObject);
				OpacityData.copyTo(Opacity);
			}

			MPlug IdPlug = Node.findPlug("particleId");
			MObject IdDataObject;
			IdPlug.getValue(IdDataObject);
			assert(IdDataObject.apiType() == MFn::kDoubleArrayData);
			MFnDoubleArrayData IdData(IdDataObject);

			Id.setLength(IdData.length());

			for (unsigned int i = 0; i < Id.length(); ++ i)
			{
				Id[i] = (int)IdData[i];
				std::cout << Id[i] << " ";
			}
		}
		else if(NodePath.hasFn(MFn::kParticle))
		{
			NodePath.extendToShape();
			MFnParticleSystem ParticleShape(NodePath.node());

			ParticleShape.position(Position);
			ParticleShape.velocity(Velocity);
			ParticleShape.radius(Radius);
			ParticleShape.rgb(Color);
			ParticleShape.opacity(Opacity);
			ParticleShape.particleIds(Id);
		}

		ParticleCount = Position.length();
		if (ParticleCount)
		{
			size_t totalDWORDCount = ParticleCount * DWORDCountOfSingleParticle;

			// Twiced size.
			Data.reset(new unsigned char[totalDWORDCount * 4 * 2]);

			float* FloatData = (float*)Data.get();
			int* IntData = (int*)Data.get();

			for(unsigned int i = 0; i < ParticleCount; ++ i)
			{
				unsigned int j = i * DWORDCountOfSingleParticle;

				FloatData[j + 0] = (float)Position[i].x;
				FloatData[j + 1] = (float)Position[i].y;
				FloatData[j + 2] = (float)Position[i].z;

				FloatData[j + 3] = (float)Velocity[i].x;
				FloatData[j + 4] = (float)Velocity[i].y;
				FloatData[j + 5] = (float)Velocity[i].z;

				FloatData[j + 6] = (float)Radius[i];

				FloatData[j + 7] = (float)Color[i].x;
				FloatData[j + 8] = (float)Color[i].y;
				FloatData[j + 9] = (float)Color[i].z;

				FloatData[j + 10] = (float)Opacity[i];

				IntData[j + 11] = Id[i];
			}
		}

		Position.clear();
		Position.setLength(0);

		Velocity.clear();
		Velocity.setLength(0);

		Radius.clear();
		Radius.setLength(0);

		Color.clear();
		Color.setLength(0);

		Opacity.clear();
		Opacity.setLength(0);

		Id.clear();
		Id.setLength(0);

		return true;
	}catch(const std::exception& e)
	{
		MGlobal::displayError(e.what());

		return false;
	}catch(...)
	{
		return false;
	}
}

bool PRTTranslator::Save(const MString& Path)
{
	try
	{
		boost::shared_array<unsigned char> Data;
		long long ParticleCount = 0;

		if (!Extract(Data, ParticleCount))
		{
			return false;
		}

		// Prepare PRT data.
		enum PRTDataType
		{
			UNKNOWN = -1,
			INT16   = 0,
			INT32   = 1,
			INT64   = 2,
			FLOAT16 = 3,
			FLOAT32 = 4,
			FLOAT64 = 5,
			UINT32  = 7,
			UINT64  = 8,
			INT8    = 9,
			INT9    = 10
		};

		struct PRTHeader
		{
			char MagicNumber[8];
			int Len;
			char Signature[32];
			int Version;
			long long Count;
		};

		PRTHeader H;
		assert(sizeof(PRTHeader) == 56);
		memset(&H, 0, sizeof(PRTHeader));
		H.MagicNumber[0] = 192;
		H.MagicNumber[1] = 'P';
		H.MagicNumber[2] = 'R';
		H.MagicNumber[3] = 'T';
		H.MagicNumber[4] = '\r';
		H.MagicNumber[5] = '\n';
		H.MagicNumber[6] = 26;
		H.MagicNumber[7] = '\n';
		H.Len = 56;
		sprintf(H.Signature,"Extensible Particle Format");
		H.Version = 1;
		H.Count = ParticleCount;

		struct PRTChannel
		{
			char Name[32];
			int DataType;
			int Arity;
			int Offset;
		};

		static const int ChannelCount = 6;

		PRTChannel Ch[ChannelCount];
		memset(Ch, 0, sizeof(PRTChannel) * ChannelCount);

		sprintf(Ch[0].Name, "Position");
		Ch[0].DataType = FLOAT32;
		Ch[0].Arity = 3;
		Ch[0].Offset = 0;

		sprintf(Ch[1].Name, "Velocity");
		Ch[1].DataType = FLOAT32;
		Ch[1].Arity = 3;
		Ch[1].Offset = 12;

		sprintf(Ch[2].Name, "Radius");
		Ch[2].DataType = FLOAT32;
		Ch[2].Arity = 1;
		Ch[2].Offset = 24;

		sprintf(Ch[3].Name, "Color");
		Ch[3].DataType = FLOAT32;
		Ch[3].Arity = 3;
		Ch[3].Offset = 28;

		sprintf(Ch[4].Name, "Opacity");
		Ch[4].DataType = FLOAT32;
		Ch[4].Arity = 1;
		Ch[4].Offset = 40;

		sprintf(Ch[5].Name, "ID");
		Ch[5].DataType = INT32;
		Ch[5].Arity = 1;
		Ch[5].Offset = 44;

		// Ready to write out.
		FILE* FP = fopen(Path.asChar(), "wb");
		if (!FP)
		{
			return false;
		}

		fwrite(&H, sizeof(PRTHeader), 1, FP);

		int RESERVED = 4;
		fwrite(&RESERVED, sizeof(int), 1, FP);
		
		fwrite(&ChannelCount, sizeof(int), 1, FP);

		int ChLength = 44;
		fwrite(&ChLength, sizeof(int), 1, FP);

		fwrite(&Ch, sizeof(PRTChannel), ChannelCount, FP);

		if (ParticleCount)
		{
			//Compress data
			const uLongf totalSize = ParticleCount * DWORDCountOfSingleParticle * 4;
			
			uLongf ZSrcBufLen = totalSize;
			Bytef* ZSrcBuf = (Bytef*)Data.get();

			uLongf ZDstBufLen = ZSrcBufLen;
			Bytef* ZDstBuf = ZSrcBuf + ZSrcBufLen;
			int ZR = compress2(ZDstBuf, &ZDstBufLen, ZSrcBuf, ZSrcBufLen, Z_BEST_COMPRESSION);
			if (ZR != Z_OK)
			{
				fclose(FP);
				remove(Path.asChar());
				return false;
			}
			fwrite(ZDstBuf, ZDstBufLen, 1, FP);
		}

		fflush(FP);
		fclose(FP);

		return true;
	}
	catch(const std::exception& e)
	{
		assert(0);
		std::cerr << e.what() << std::endl;
		return false;
	}
	catch(...)
	{
		assert(0);
		return false;
	}
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#ifdef _WIN32
	#define EXPORT __declspec(dllexport)
#elif __GNUC__ >= 4
	#define EXPORT __attribute__ ((visibility("default")))
#else
	#define EXPORT
#endif

EXPORT
MStatus initializePlugin(MObject Object)
{
	MFnPlugin Plugin(Object, "Bo Zhou", "2.0", "10.0");
	MStatus R = Plugin.registerFileTranslator("PRT", "", PRTTranslator::Creator);
	if (!R) 
	{
		R.perror("register PRTTranslator");
	}
	return R;
}

EXPORT
MStatus uninitializePlugin(MObject Object)
{
	MFnPlugin Plugin(Object);
	MStatus R = Plugin.deregisterFileTranslator("PRT");
	if (!R)
	{
		R.perror("deregister PRTTranslator");
	}
	return R;
}
